package main

import (
	"errors"
	"fmt"
	"github.com/spf13/pflag"
	"math"
	"nprobe/bootstrap"
	"nprobe/core"
	"nprobe/core/manager"
	"nprobe/utils/nets"
	"nprobe/utils/routines"
	"os"
	"strconv"
	"strings"
	"time"
)

var help = pflag.BoolP("help", "h", false, "help info")
var udp = pflag.BoolP("udp", "U", false, "use udp datagram instead of icmp echo")
var tcpsyn = pflag.BoolP("tcpsyn", "S", false, "use tcp syn packet instead of icmp echo")
var tcp = pflag.BoolP("tcp", "T", false, "use tcp packet instead of icmp echo")
var server = pflag.BoolP("listen", "l", false, "work as a server")
var sources = pflag.StringP("sources", "s", "", "source address, while use icmp echo, use multi sources to cover ecmp paths")
var srcPorts = pflag.String("sport", "63000-63999", "source port or port range, while use udp or tcp, use multi sport and dport to cover ecmp paths")
var dstPorts = pflag.String("dport", "53", "destination port or port range, while use udp or tcp, use multi sport and dport to cover ecmp paths")
var tos = pflag.Uint8("tos", 0, "tos")
var ttl = pflag.Uint8("ttl", 64, "ttl")
var length = pflag.Int("length", 64, "payload data length")
var timeout = pflag.Int64P("timeout", "t", 1000, "message response timeout time, measured in milliseconds")
var count = pflag.IntP("count", "c", math.MaxInt, "repeat count of probe flow")

var pps = pflag.Int("pps", 10, "packets count per second of a target")
var span = pflag.Int64("span", 1000, "time duration for statistical, measured in milliseconds")
var bitFlip = pflag.Bool("bit-flip", false, "check packet bit flip events")
var withEchoServer = pflag.Bool("with-echo-server", false, "whether udp probe working with a echo server")

func main() {
	if helpFunc() {
		return
	}

	err := bootstrap.InitEnv()
	if err != nil {
		panic(err)
	}

	task := buildProbeTask()
	if task == nil {
		return
	}

	agent, err := manager.NewAgent(task)
	if err != nil {
		return
	}

	err = agent.StartAgent()
	if err != nil {
		agent.StopAgent()
	}

	routines.WaitExit()
	time.Sleep(time.Second)
	agent.StopAgent()
}

func buildProbeTask() *core.ProbeTask {
	var err error
	proto := determineProto()

	var dstIPs []string
	var srcIPs []string
	if pflag.NArg() > 0 {
		dstIPs, err = parseTarget(pflag.Args()[0])
		if err != nil {
			return nil
		}
	}

	srcIPs, err = parseSources(*sources, dstIPs, proto)
	if err != nil {
		return nil
	}

	srcPortList := []uint16{0, 0}
	dstPortList := []uint16{0, 0}
	if proto == "udp" || proto == "tcp" || proto == "tcpsyn" {
		srcPortList, err = calcPorts(*srcPorts)
		if err != nil {
			return nil
		}
		dstPortList, err = calcPorts(*dstPorts)
		if err != nil {
			return nil
		}
	}

	if *length < 22 {
		*length = 22
	}

	role := "client"
	if *server {
		role = "server"
	}

	if *ttl <= 0 || *timeout <= 0 || *count <= 0 {
		return nil
	}

	return &core.ProbeTask{
		Role:           role,
		DstIPs:         dstIPs,
		SrcIPs:         srcIPs,
		DstPortMin:     dstPortList[0],
		DstPortMax:     dstPortList[1],
		SrcPortMin:     srcPortList[0],
		SrcPortMax:     srcPortList[1],
		Proto:          proto,
		IPVer:          4,
		TTL:            *ttl,
		MsgLen:         *length,
		TOS:            *tos,
		TimeoutMs:      *timeout,
		Count:          *count,
		PPS:            *pps,
		SpanMs:         *span,
		BitFlip:        *bitFlip,
		WithEchoServer: *withEchoServer,
	}
}

func determineProto() string {
	proto := "icmp"
	if *udp {
		proto = "udp"
	} else if *tcp {
		proto = "tcp"
	} else if *tcpsyn {
		proto = "tcpsyn"
	}

	return proto
}

func calcPorts(ports string) ([]uint16, error) {

	var minPort uint16
	var maxPort uint16
	portRange := strings.Split(ports, "-")
	if len(portRange) == 1 {
		port, err := strconv.ParseUint(portRange[0], 10, 16)
		if err != nil {
			fmt.Println("ntrace: unsupported port value")
			return nil, err
		}
		minPort = uint16(port)
		maxPort = uint16(port)
	} else if len(portRange) == 2 {
		low, err := strconv.ParseUint(portRange[0], 10, 16)
		if err != nil {
			fmt.Println("ntrace: unsupported port value")
			return nil, err
		}
		high, _ := strconv.ParseUint(portRange[1], 10, 16)
		if err != nil {
			fmt.Println("ntrace: unsupported port value")
			return nil, err
		}
		if low > high {
			fmt.Println("ntrace: unsupported port value")
			return nil, err
		}
		minPort = uint16(low)
		maxPort = uint16(high)
	} else {
		fmt.Println("ntrace: unsupported port value")
		return nil, errors.New("ntrace: unsupported port value")
	}

	parsedPorts := []uint16{minPort, maxPort}
	return parsedPorts, nil
}

func parseTarget(targets string) ([]string, error) {
	var dstIPs []string
	segs := strings.Split(targets, ",")
	for _, seg := range segs {
		parsedDstIPs, err := nets.ParseIPAddr(seg)
		if err != nil {
			continue
		}
		dstIPs = append(dstIPs, parsedDstIPs...)
	}

	if len(dstIPs) == 0 {
		fmt.Println("nprobe: please provide target")
		return nil, errors.New("nprobe: please provide target")
	}

	return dstIPs, nil
}

func parseSources(sources string, dstIPs []string, role string) ([]string, error) {
	var srcIPs []string
	segs := strings.Split(sources, ",")
	for _, seg := range segs {
		parsedSources, err := nets.ParseIPAddr(seg)
		if err != nil {
			continue
		}
		srcIPs = append(srcIPs, parsedSources...)
	}

	if len(srcIPs) == 0 {
		for _, dstIP := range dstIPs {
			srcIP, err := nets.LookupLocalAddrByRoute(dstIP)
			if err != nil {
				continue
			}
			srcIPs = append(srcIPs, srcIP.To4().String())
		}
	}

	if len(srcIPs) == 0 && role == "client" {
		fmt.Println("nprobe: can not find any valid source ip")
		return nil, errors.New("nprobe: can not find any valid source ip")
	}

	return srcIPs, nil
}

func helpFunc() bool {
	pflag.Usage = func() {
		fmt.Printf("Usage of %s:\n", os.Args[0])
		fmt.Printf("%s [-options] target\n", os.Args[0])
		pflag.PrintDefaults()
	}

	pflag.Parse()
	if *help {
		pflag.Usage()
		return true
	}

	if pflag.NArg() == 0 && !(*server) {
		pflag.Usage()
		return true
	}

	return false
}
